import torch.nn as nn
'''
  MCNN网络
  输入 x1（1,n,n）  
  输出（1,n，n）
  示例：
    x1 = torch.rand([100, 1, 69, 69])
    model = MCNN(num_convolutions=3, kernel_size=3, scale=2)
    y = model(x1)
    print(y.shape)
  注意：
    kernel_size必须为单数，否则无法保证每次卷积后，size不变
    O = （I - k + 2p）/S + 1
'''

class MCNN(nn.Module):
    def __init__(self, num_convolutions, kernel_size, scale = 2):
        super(MCNN, self).__init__()
        self.layers = nn.ModuleList()
        for i in range(num_convolutions):
            self.layers.append(nn.Conv2d(scale ** i, scale ** (i+1), kernel_size=kernel_size, padding=(kernel_size-1)//2, stride=1))
    def forward(self, x):
        for layer in self.layers:
            x = layer(x)
        x = x.sum(dim=1, keepdim=True)
        return x